# 导入操作系统相关的功能模块
import os

# 导入 PyTorch 深度学习框架
import torch
# 从 PyTorch 中导入神经网络模块
from torch import nn
# 从 PyTorch 中导入 AdamW 优化器
from torch.optim import AdamW
# 从 PyTorch 中导入学习率调度器，包括余弦退火调度器和多步调度器
from torch.optim.lr_scheduler import CosineAnnealingLR, MultiStepLR
# 从 PyTorch 中导入数据加载器
from torch.utils.data import DataLoader

# 从自定义的 dataset 模块中导入 SwitchDataset 类
from dataset import SwitchDataset
# 从自定义的 model 模块中导入 SwitchModel 类
from model import SwitchModel


def main(data_root, output_dir, backbone, total_epochs):
    """
    主训练函数，用于训练一个分类模型。

    参数:
    data_root (str): 数据集的根目录。
    output_dir (str): 模型保存的目录。
    """
    # 检查是否有可用的 CUDA 设备，如果有则使用 GPU，否则使用 CPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # 创建 SwitchDataset 实例，传入数据集根目录和图像大小
    dataset = SwitchDataset(data_root, size=224)
    # 获取数据集中的类别名称列表
    class_names = dataset.get_class_names()
    class_to_idx = dataset.get_class_to_idx()
    print(class_to_idx)
    # 创建数据加载器，用于批量加载数据
    dataloader = DataLoader(
        dataset=dataset,  # 指定数据集
        batch_size=16,  # 每个批次的样本数量
        shuffle=True,  # 是否在每个 epoch 开始时打乱数据
        num_workers=4,  # 用于数据加载的子进程数量
        pin_memory=True,  # 是否将数据加载到固定内存中，以加快数据传输速度
        drop_last=True,  # 是否丢弃最后一个不完整的批次
    )
    
    # 创建 SwitchModel 实例，指定骨干网络、类别数量和是否使用预训练权重
    model = SwitchModel(backbone=backbone, num_classes=len(class_names), pretrained=True)
    # 将模型移动到指定的设备（GPU 或 CPU）上
    model.to(device)
    
    # 创建 AdamW 优化器，传入模型的参数和学习率
    optimizer = AdamW(model.parameters(), lr=1e-4)
    # 创建余弦退火学习率调度器，传入优化器、最大迭代次数和最小学习率
    scheduler = CosineAnnealingLR(optimizer, T_max=total_epochs, eta_min=1e-6)
    
    # 定义交叉熵损失函数，用于分类任务
    criterion = nn.CrossEntropyLoss()
    
    # 开始训练循环，遍历每个 epoch
    for epoch in range(total_epochs):
        # 将模型设置为训练模式
        model.train()
        # 初始化总损失为 0
        total_loss = 0
        # 初始化正确预测的样本数量为 0
        total_correct = 0
        # 初始化总样本数量为 0
        total_num = 0
        # 遍历数据加载器中的每个批次
        for images, labels in dataloader:
            # 将图像数据移动到指定的设备上
            images = images.to(device)
            # 将标签数据移动到指定的设备上
            labels = labels.to(device)

            # 清零优化器中的梯度信息
            optimizer.zero_grad()
            # 前向传播，通过模型得到输出
            outputs = model(images)
            # 计算损失，使用交叉熵损失函数
            loss = criterion(outputs, labels)
            # 反向传播，计算梯度
            loss.backward()
            # 更新模型的参数
            optimizer.step()

            # 累加当前批次的损失到总损失中
            total_loss += loss.item()
            # 累加当前批次中正确预测的样本数量到总正确数量中
            total_correct += (outputs.argmax(1) == labels).sum().item()
            # 累加当前批次的样本数量到总样本数量中
            total_num += labels.shape[0]
        # 更新学习率调度器
        scheduler.step()

        # 打印当前 epoch 的训练损失和准确率
        print(f"Epoch {epoch + 1}/{total_epochs}, loss: {total_loss / total_num}, acc: {total_correct / total_num}")
    
    # 确保输出目录存在，如果不存在则创建
    os.makedirs(output_dir, exist_ok=True)
    # 将模型设置为评估模式
    model.eval()
    # 保存模型的状态字典到指定的输出目录
    torch.save(model.state_dict(), os.path.join(output_dir, f"model_{backbone}.pth"))
    # 打印模型保存成功的消息
    print("Model saved!")
    
    # 打印训练完成的消息
    print("Finish!")

# 当脚本作为主程序运行时执行以下代码
if __name__ == "__main__":
    # 定义数据集的根目录
    data_root = "data/dataset"
    # 定义模型保存的输出目录
    output_dir = "output"
    # 定义骨干网络的类型
    backbone = "mobilenet" # 可选 mobilenet, shufflenet
    # 定义训练的总轮数
    # total_epochs = 50
    if backbone == "mobilenet":
        total_epochs = 20
    else:
        total_epochs = 50

    # 调用主训练函数
    main(
        data_root=data_root,
        output_dir=output_dir,
        backbone=backbone,
        total_epochs=total_epochs,
    )